import numpy as np
import win32gui
import win32ui
from ctypes import windll
from PIL import Image


def screenshot_py32(window_title: str):
    # 根据窗口标题找到目标窗口句柄
    hwnd = win32gui.FindWindow(None, window_title)
    if not hwnd:
        return None

    # 获取窗口的大小
    left, top, right, bottom = win32gui.GetWindowRect(hwnd)
    width = right - left
    height = bottom - top

    # 创建设备上下文（DC）
    hwindc = win32gui.GetWindowDC(hwnd)
    srcdc = win32ui.CreateDCFromHandle(hwindc)
    memdc = srcdc.CreateCompatibleDC()

    # 创建位图对象
    bmp = win32ui.CreateBitmap()
    bmp.CreateCompatibleBitmap(srcdc, width, height)
    memdc.SelectObject(bmp)

    # 将窗口内容复制到位图
    result = windll.user32.PrintWindow(hwnd, memdc.GetSafeHdc(), 2)
    if result == 1:
        # 将位图转换为 PIL 图像对象
        bmpinfo = bmp.GetInfo()
        bmpstr = bmp.GetBitmapBits(True)
        # img = Image.frombuffer(
        #     'RGB',
        #     (bmpinfo['bmWidth'], bmpinfo['bmHeight']),
        #     bmpstr, 'raw', 'BGRX', 0, 1
        # )
        img = np.frombuffer(bmpstr, dtype=np.uint8).reshape((bmpinfo["bmHeight"], bmpinfo["bmWidth"], 4))
        img = np.ascontiguousarray(img)[..., :-1]

        # 清理资源
        win32gui.DeleteObject(bmp.GetHandle())
        memdc.DeleteDC()
        srcdc.DeleteDC()
        win32gui.ReleaseDC(hwnd, hwindc)

        return img
    else:
        return None

if __name__ == '__main__':
    window_title = "管理员: C:\Windows\system32\cmd.exe"
    screenshot = screenshot_py32(window_title)
    if screenshot:
        screenshot.save('screenshot.png')
        print("Screenshot saved.")
    else:
        print("Failed to capture window.")